# 18/100 矩阵-矩阵置零
# leetcode第73题: https://leetcode.cn/problems/set-matrix-zeroes/description/?envType=study-plan-v2&envId=top-100-liked
# Date: 2024/11/22
from leetcode import test


def setZeroes(matrix: list[list[int]]) -> list[list[int]]:
    row = set()
    col = set()
    for i in range(len(matrix)):
        for j in range(len(matrix[0])):
            if matrix[i][j] == 0:
                row.add(i)
                col.add(j)

    for i in range(len(matrix)):
        for j in range(len(matrix[0])):
            if i in row or j in col:
                matrix[i][j] = 0

    return matrix


def setZeroes_inplace(matrix: list[list[int]]) -> list[list[int]]:
    """使用原地标记的方法
    使用两个标记变量的方法: 我们可以用矩阵的第一行和第一列代替方法一中的两个标记数组，以达到 O(1) 的额外空间。
    但这样会导致原数组的第一行和第一列被修改，无法记录它们是否原本包含 0。
    因此我们需要额外使用两个标记变量分别记录第一行和第一列是否原本包含 0。
    在实际代码中，我们首先预处理出两个标记变量，接着使用其他行与列去处理第一行与第一列，然后反过来使用第一行与第一列去更新其他行与列，
    最后使用两个标记变量更新第一行与第一列即可。
    """
    m, n = len(matrix), len(matrix[0])
    col0 = any(matrix[i][0] == 0 for i in range(m))
    row0 = any(matrix[0][j] == 0 for j in range(n))

    for i in range(1, m):
        for j in range(1, n):
            if matrix[i][j] == 0:
                matrix[i][0] = matrix[0][j] = 0

    for i in range(1, m):
        for j in range(1, n):
            if matrix[i][0] == 0 or matrix[0][j] == 0:
                matrix[i][j] = 0

    if col0:
        for i in range(m):
            matrix[i][0] = 0

    if row0:
        for j in range(n):
            matrix[0][j] = 0
    return matrix


if __name__ == '__main__':
    inp = [{"matrix": [[1, 1, 1], [1, 0, 1], [1, 1, 1]]}, {"matrix": [[0, 1, 2, 0], [3, 4, 5, 2], [1, 3, 1, 5]]}, ]
    out = [[[1, 0, 1], [0, 0, 0], [1, 0, 1]], [[0, 0, 0, 0], [0, 4, 5, 0], [0, 3, 1, 0]]]
    test.test_function(setZeroes, inp, out)
    test.test_function(setZeroes_inplace, inp, out)
